import os
import json
import argparse
import sys 
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

sys.path.append("../../")
from package.utility import ensure_exists
from package.utility import find_targets
sys.path.append("..")
from baseplot import casename2number
from baseplot import create_subplot
from baseplot import mypaper_subplot
from baseplot import close_sublot_levelplot
from baseplot import mypaper_subplot_level
from baseplot import mypaper_subplot_level_axset

def get_runcases(topdir,jsonname="*.json"):
    """Get runcase under topdir."""
    topdir = os.path.abspath(topdir)
    ensure_exists(topdir)

    print(f">> scan runcase under {topdir}") 
    runcases = find_targets(topdir, jsonname, suffix="-xtype f") 
    runcase_num = len(runcases)
    print(f">> {runcase_num} runcases found!")

    return runcases

def getjsondata(casepath):
    result_list = []
    basename = os.path.basename(casepath)
    programname = basename.split("-")[0]
    with open(casepath, "r") as fo:
        text = fo.read()
        jsondata = json.loads(text)
        print(f"Load {basename} data.")
    for i in jsondata:
        timedic = {}
        timedic["np"] = int(i.split("_")[-1])
        for j in jsondata[i]:
            if j == "FindIndependentset":
                timedic[programname] = float(jsondata[i][j]["tavg"])
        result_list.append(timedic)
    return result_list

def mycsv(savedir,basename,dict):
    df = pd.DataFrame(dict)
    df = pd.DataFrame(dict)
    df = df.sort_index(axis=1)
    df = df.sort_values(by='np', ascending=True)
    df.set_index(["np"], inplace=True)
    
    # print(df)
    # exit(0)
    writer = pd.ExcelWriter(os.path.join(savedir, basename+".xlsx"))
    df.to_excel(writer)
    writer.save()

def main():
    dir = "/home/mrz/data/new_allreduce/ok_data/230425T2_useinamg/tianhe3_testpdl"
    runcases = get_runcases(dir,"*.matrix-database.json")
    savedir = "/home/mrz/data/new_allreduce/ok_data/dataprocess0510"

    # print(runcases)
    for i in runcases:
        basename = os.path.basename(i)
        print(f"read {i}\n")
        datadict = getjsondata(i)
        mycsv(savedir,basename,datadict)
    # # print(list(datadict["jxpamg_solver04"]["np"]))
    # myplot(args.savedir,datadict)
    # # mycsv(args.savedir,datadict)

if __name__ == "__main__":
    main()